/*
 * 内存池实现
 *
 * Author: couryzh
 * Email:  couryzh@aliyun.com
 *
 */
#include <string.h>
#include <stdio.h>
#include <stddef.h>
#include <pthread.h>
#include "mempool.h"
#include "log.h"


struct mem_node;

struct mem_block{
    size_t mem_size;
    size_t num_node;
    struct mem_block *next;
    char mem[0];
};

typedef struct mem_block mem_block_chain_t;

struct mem_group
{
    size_t size;
    struct mem_pool *pool;
    mem_block_chain_t *block_chain;
    struct mem_node *head, *tail;
};

// 内存头，从系统申请时加上
typedef union
{
    struct mem_group *group;  // 分配出去时，填充所有blk
    struct mem_node *next; // 归还后，串入free list
} union_pad_t;

// 内存单元
struct mem_node
{
    union_pad_t u; // 内存头
    char mem[0];   // 实际使用内存，大小对应所属mem_block.size
};

// 内存池
struct mem_pool
{
    size_t max_alloc;  // 允许从系统申请数量
    size_t alloced; // 已申请数量
    int concurrency;
    pthread_mutex_t lock;
    int num_group;     //
    struct mem_group group[0];
};


// 对齐
#define ALIGN(n, s)     (((n+s-1)/s) * s)

// 每个block分配字节数
#define MEMORY_SIZE  (4096)

// 每个block大小
#define BLOCK_SIZE      (sizeof(struct mem_block) + MEMORY_SIZE)


static int expand_group(struct mem_group *group)
{
    int i, offset;
    struct mem_node *node = NULL;
    struct mem_block *block = NULL, **blkp = NULL;

    log_debug("alloc block %d bytes\n", BLOCK_SIZE);
    block = malloc(BLOCK_SIZE);
    if (!block) {
        log_err("failed to expand block\n");
        return 0;
    }
    block->mem_size = MEMORY_SIZE;
    block->next = NULL;

    for (i=0; ; i++) {
        offset = ALIGN((sizeof(struct mem_node) + group->size)*i, 4);
        if (offset >= block->mem_size) break;
        node = (struct mem_node *)(((char *)block->mem) + offset);
        node->u.next = NULL;
        if (!group->head)
            group->head = group->tail = node;
        else {
            group->tail->u.next = node;
            group->tail = node;
        }
    }
    block->num_node = i;

    // 新block挂到block_chain最后
    for (blkp=&group->block_chain; *blkp; blkp=&(*blkp)->next) {}
    *blkp = block;

    group->pool->alloced += block->mem_size;
    log_debug("alloced size: %d\n", group->pool->alloced);

    return 1;
}

static inline void pool_lock_init(mem_pool_t *p, int concurrency)
{
    p->concurrency = concurrency;
    if (p->concurrency) {
        pthread_mutex_init(&p->lock, NULL);
    }
}
static inline void pool_lock_destroy(mem_pool_t *p)
{
    if (p->concurrency)
        pthread_mutex_destroy(&p->lock);
}

static inline void pool_lock(mem_pool_t *p)
{
    if (p->concurrency)
        pthread_mutex_lock(&p->lock);
}

static inline void pool_unlock(mem_pool_t *p)
{
    if (p->concurrency)
        pthread_mutex_unlock(&p->lock);
}

int mempool_init(mem_pool_t **pool, int max_alloc, int num_level, int *levels, int concurrency)
{
    int i, j;
    mem_pool_t *p = NULL;
    struct mem_group *group = NULL;

    p = calloc(1, sizeof(mem_pool_t) + num_level*sizeof(struct mem_group));
    if (!p) {
        log_err("failed to alloc memory for pool\n");
        return ERR_NO_MEM;
    }

    for (i=0; i<num_level; i++) {
        group = &p->group[i];
        group->pool = p;
        group->size = levels[i];
        if (!expand_group(group))
			goto err_return;
    }
    p->max_alloc = max_alloc;
    p->num_group = num_level;

    pool_lock_init(p, concurrency);

    *pool = p;
    return 0;

err_return:
    for (j=0; j<i; j++) {
        free(group->block_chain);
    }
    free(p);
    return ERR_NO_MEM;
}


void *mempool_alloc(mem_pool_t *pool, size_t siz)
{
    int i;
    struct mem_group *group = NULL;
    struct mem_node *node = NULL;

    // 找到siz对应的合适group
    for (i=0; i<pool->num_group; i++) {
        if (pool->group[i].size >= siz) {
            group = &pool->group[i];
            break;
        }
    }
    if (!group) {
        log_err("alloc size %d overrange\n", siz);
        return NULL;
    }

    pool_lock(pool);
    // group是否还有free node
    if (!group->head) {
        // 无free node,但已达到max_alloc
        if (pool->alloced + BLOCK_SIZE >= pool->max_alloc) {
            log_warn("alloced memory will over max alloc %d\n", pool->max_alloc);
            pool_unlock(pool);
            return NULL;
        }
        // 未到max_alloc, 进行扩展
        if (!expand_group(group)) return NULL;
    }
    // 从head取
    node = group->head;
    group->head = group->head->u.next;
    node->u.group = group;
    if (group->tail == node)
        group->tail = NULL;
    pool_unlock(pool);
    // 使用mem
    return node->mem;
}

void mempool_free(mem_pool_t *pool, void *p)
{
    struct mem_node *node = NULL;
    struct mem_group *group = NULL;

    // 定位到mem_node容器
    node = (struct mem_node *)(((char *)p) - offsetof(struct mem_node, mem));
    group = node->u.group;
    // 挂到tail
    pool_lock(pool);
    node->u.next = NULL;
    if (!group->tail) {
        group->head = group->tail = node;
    }
    else {
        group->tail->u.next = node;
        group->tail = node;
    }
    pool_unlock(pool);
}

void mempool_dump(char *desc, mem_pool_t *p)
{
    int i;
    struct mem_block *blk = NULL;
    struct mem_node *node = NULL;

    printf("%s\n", desc);
    printf("max_alloc: %lu, "
            "alloced: %lu, "
            "concurrency: %s, "
            "num group:%d\n"
            "--------------------------\n", p->max_alloc, p->alloced, p->concurrency ? "yes": "no", p->num_group);
    for (i=0; i<p->num_group; i++) {
        printf("size: %lu\n", p->group[i].size);
        printf("free list: ");
        for (node = p->group[i].head; node; node=node->u.next) {
            printf("%p ", node->mem);
        }
        printf("\n---------------------\n");
        printf("block(num): ");
        for (blk=p->group[i].block_chain; blk; blk=blk->next) {
            printf("%lu ", blk->num_node);
        }
        printf("\n-----------------------------------\n");
    }
}

void dump_mem_node(void *p)
{
    struct mem_node *node = NULL;
    node = (struct mem_node *)(((char *)p) - offsetof(struct mem_node, mem));
    printf("memory: %p, group-size: %lu\n", node->mem, node->u.group->size);
}

void mempool_destroy(mem_pool_t **pool)
{
    int i;
    struct mem_block *blk = NULL, *next = NULL;
    mem_pool_t *p = *pool;

    // 遍历释放每个block
    for (i=0; i<p->num_group; i++) {
        for (blk=p->group[i].block_chain; blk; blk=next) {
            next = blk->next;
            free(blk);
        }
    }

    pool_lock_destroy(p);
    p->alloced = 0;
    free(p);
    *pool = NULL;
}

